import torch
import torch.nn as nn
from EnergyConfig import EnergyConfig


class EnergyModel(nn.Module):

    def __init__(self):
        super(EnergyModel, self).__init__()
        self.hidden_size = EnergyConfig.hidden_size
        self.num_layers = EnergyConfig.num_layers

        self.rnn = nn.LSTM(
            input_size=EnergyConfig.n_features,
            hidden_size=EnergyConfig.hidden_size,
            num_layers=EnergyConfig.num_layers,
            batch_first=True,
            dropout=EnergyConfig.dropout_prob if EnergyConfig.num_layers > 1 else 0,
        )

        self.dropout = nn.Dropout(EnergyConfig.dropout_prob)
        self.fc = nn.Linear(EnergyConfig.hidden_size, EnergyConfig.num_classes)

    def forward(self, x):
        h0 = torch.zeros(self.num_layers, x.size(0), self.hidden_size).to(x.device)
        c0 = torch.zeros(self.num_layers, x.size(0), self.hidden_size).to(x.device)
        out, _ = self.rnn(x, (h0, c0))
        out = out[:, -1, :]
        out = self.dropout(out)
        out = self.fc(out)
        return out.squeeze()
